/*
 * GraphState2.cpp
 *
 *  Created on: Jun 28, 2011
 *      Author: pickrell
 */

#include "Tree.h"
using namespace std;

Tree::Tree(){
}

Tree::Tree(string instring){
	//cout << "going to set graph0?\n"; cout.flush();
	tree = new PopGraph();
	//cout << "going to set graph?\n"; cout.flush();
	set_graph(instring);
}

void Tree::set_graph(string newick){
	//cout << "calling set_graph\n"; cout.flush();
	npop = 0;
	//cout << "setting g\n"; cout.flush();
	set_graph_from_string(newick);
	//cout << "setting c\n"; cout.flush();
	set_coaltimes();
	//cout << "setting h\n"; cout.flush();
	set_heights();
	//cout << "done\n"; cout.flush();
}

void Tree::set_graph(string vfile, string efile){
	igzstream vin(vfile.c_str());
	tree->g.clear();
    vector<string> line;
    struct stat stFileInfo;
    int intStat;

    tree->popnames.clear();

    intStat = stat(vfile.c_str(), &stFileInfo);
    if (intStat !=0){
            std::cerr<< "ERROR: cannot open file " << vfile << "\n";
            exit(1);
    }
    string st;
    while(getline(vin, st)){
    	string buf;
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		int index = atoi(line[0].c_str());
		if (index >= tree->indexcounter) tree->indexcounter = index+1;
		string name = line[1];
		string root = line[2];
		string mig = line[3];
		string tip = line[4];
		Graph::vertex_descriptor v= add_vertex(tree->g);
		tree->g[v].index = index;
		tree->g[v].name = name;
		tree->g[v].mig_frac = 1;
		//tree->g[v].popid = 0;
		if (root == "ROOT") {
			tree->g[v].is_root = true;
			tree->root = v;
		}
		else tree->g[v].is_root = false;
		if (mig == "MIG") {
			tree->g[v].is_mig = true;
			tree->g[v].mig_frac = 0.5;
		}
		else tree->g[v].is_mig = false;
		if (tip == "TIP") {
			tree->g[v].is_tip = true;
			npop++;
			tree->popnames.push_back(name);
		}
		else tree->g[v].is_tip = false;
		tree->g[v].rev = false;
     }

	igzstream ein(efile.c_str());
	map<int, Graph::vertex_descriptor> i2v = tree->index2vertex();
    intStat = stat(efile.c_str(), &stFileInfo);
    if (intStat !=0){
            std::cerr<< "ERROR: cannot open file " << efile << "\n";
            exit(1);
    }
    while(getline(ein, st)){
    	string buf;
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		int index1 = atoi(line[0].c_str());
		int index2 = atoi(line[1].c_str());
		float len = atof(line[2].c_str());
		float w = atof(line[3].c_str());
		string mig= line[4];
		Graph::edge_descriptor e = add_edge( i2v[index1], i2v[index2], tree->g).first;
		Graph::vertex_descriptor t = i2v[index2];
		if (tree->g[t].is_mig) w = 1;
		tree->g[e].len = len;
		tree->g[e].weight = w;
		if (mig == "MIG") {
			//cout << "here\n";
			tree->g[e].is_mig = true;
			//float mig_frac = atof(line[5].c_str());
			//tree->set_mig_frac(e, mig_frac);
			//cout << "not here\n";
		}
		else tree->g[e].is_mig = false;
    }
    ein.close();
    igzstream ein2(efile.c_str());
    while(getline(ein2, st)){
     	string buf;
 		stringstream ss(st);
 		line.clear();
 		while (ss>> buf){
 			line.push_back(buf);
 		}
 		int index1 = atoi(line[0].c_str());
 		int index2 = atoi(line[1].c_str());
 		string mig= line[4];
 		Graph::edge_descriptor e = edge( i2v[index1], i2v[index2], tree->g).first;
 		if (mig == "MIG") {
 			//cout << "here\n";
 			//tree->g[e].is_mig = true;
 			float mig_frac = atof(line[5].c_str());
 			//tree->set_mig_frac(e, mig_frac);
 			//cout << "not here\n";
 		}
     }
    //tree->print();



}




void Tree::set_graph_from_string(string newick){
	npop = 0;
	tree->set_graph(newick);
	map<string, Graph::vertex_descriptor> tips = tree->get_tips(tree->root);
	vector<string> allpopnames;
	for (map<string, Graph::vertex_descriptor>::iterator it = tips.begin(); it != tips.end(); it++){
		allpopnames.push_back(it->first);

	}
    npop = allpopnames.size();

	//tree->print();

}

map<Graph::vertex_descriptor, int> Tree::get_v2index(){

	map<Graph::vertex_descriptor, int> vertex2index;

	vector<Graph::vertex_descriptor> i_nodes = tree->get_inorder_traversal(npop); //get descriptors for all the nodes
	set<Graph::vertex_descriptor> root_adj = tree->get_root_adj(); //get the ones next to the root
	vector<Graph::vertex_descriptor> i_nodes2;  //remove the ones next to the root

	int index = 0;

	for (int i = 0; i < i_nodes.size(); i++){
		if (root_adj.find(i_nodes[i]) == root_adj.end() && !tree->g[ i_nodes[i] ].is_root) {
			i_nodes2.push_back( i_nodes[i] );
			vertex2index.insert(make_pair(i_nodes[i], index));
			index++;
		}
	}

	int joint_index = i_nodes2.size(); //the index of the parameter for the sum of the two branch lengths next to the root

	for(set<Graph::vertex_descriptor>::iterator it = root_adj.begin(); it != root_adj.end(); it++) vertex2index[*it] = joint_index;
	index++;
	//get all the migration nodes
	vector<Graph::vertex_descriptor> mig_nodes;
	for(Graph::vertex_iterator it = vertices(tree->g).first; it != vertices(tree->g).second; it++){
		if ( tree->g[*it].is_mig ) {
			mig_nodes.push_back(*it);
			vertex2index[*it] = index;
			index++;
		}
	}
	return vertex2index;

}

void Tree::print_coales(){
	set_labeled_coaltimes();
	for(vector<double>::iterator it = coaltimes1.begin(); it != coaltimes1.end(); it++){
		cout << *it << " 1\n";
	}
	for(vector<double>::iterator it = coaltimes2.begin(); it != coaltimes2.end(); it++){
		cout << *it << " 2\n";
	}
	for(vector<double>::iterator it = coaltimes3.begin(); it != coaltimes3.end(); it++){
		cout << *it << " 3\n";
	}

}

void Tree::set_coaltimes(){
	coaltimes.clear();
	vector<Graph::vertex_descriptor> v= tree->get_inorder_traversal(npop);
	//cout << "here2\n"; cout.flush();
	for (vector<Graph::vertex_descriptor>::iterator it = v.begin(); it != v.end(); it++){
		if (tree->g[*it].is_tip == false){
			float tmp = tree->get_coaltime(*it);
			//cout << tmp << " tmp\n";
			coaltimes.push_back(tmp);
		}
	}
	sort(coaltimes.begin(), coaltimes.end());
}
void Tree::set_labeled_coaltimes(){
	coaltimes1.clear();
	coaltimes2.clear();
	coaltimes3.clear();
	vector<Graph::vertex_descriptor> v= tree->get_inorder_traversal(npop);
	for (vector<Graph::vertex_descriptor>::iterator it = v.begin(); it != v.end(); it++){
		if (tree->g[*it].is_tip == false){
			if (tree->g[*it].popid == 1) coaltimes1.push_back(tree->g[*it].height);
			else if (tree->g[*it].popid == 2) coaltimes2.push_back(tree->g[*it].height);
			else if (tree->g[*it].popid == 3) coaltimes3.push_back(tree->g[*it].height);
		}
	}
	sort(coaltimes1.begin(), coaltimes1.end());
	sort(coaltimes2.begin(), coaltimes2.end());
	sort(coaltimes3.begin(), coaltimes3.end());
}
double Tree::llk(vector<double>* bpoints, vector<double>* popsizes){
	double toreturn = 0;
	int i = 0; //index coaltimes
	int j = 0; //index bpoints
	int k = 0; //index popsizes
	double previous_coaltime = 0;
	double previous_bpoint = 0;
	double previous_popsize = popsizes->at(0);
	while (i < coaltimes.size()){
		double current_coaltime = coaltimes.at(i);
		double current_bpoint = bpoints->at(j);
		double current_popsize = popsizes->at(k);
		int nlin = coaltimes.size()-i+1;
		//cout << i << " "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
		double add = 0;
		double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		//cout << "C "<< C << "\n";
		while (current_bpoint < current_coaltime){



			//cout <<"here "<< i << " "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";

			//add the bit after the previous event
			if (previous_coaltime < current_bpoint && previous_coaltime > previous_bpoint)	{

				//cout << i << " here2 "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << "\n";
					//cout << i << " adding2 "<< - (current_bpoint - previous_coaltime) * C / current_popsize << "\n";
					//cout << previous_bpoint - previous_coaltime << " "<< C / previous_popsize << "\n";
					add += - (current_bpoint - previous_coaltime) * C / current_popsize;

			}
			else	{
				//cout << i << " here3 "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
				//cout << i  << " adding4 " << - (current_bpoint - previous_bpoint)* C / current_popsize <<"\n";
				add += - (current_bpoint - previous_bpoint)* C / current_popsize;
			}
			//increment
			j++;
			k++;
			previous_bpoint = current_bpoint;
			previous_popsize = current_popsize;
			current_bpoint = bpoints->at(j);
			current_popsize = popsizes->at(k);

		}
		if (previous_coaltime > previous_bpoint){
			//cout << "here4 "<< current_coaltime << " "<< previous_coaltime << " "<< C << " "<< current_popsize << " \n";
			//cout << i  << " adding5 " <<-(current_coaltime - previous_coaltime)*  C/ current_popsize << "\n";
			add += -(current_coaltime - previous_coaltime)*  C/ current_popsize;
		}
		else {
			//cout << i << " adding6 " << - (current_coaltime - previous_bpoint) * C /current_popsize << "\n";
			add += - (current_coaltime - previous_bpoint) * C /current_popsize;
		}
		//cout << i << " adding7 " << log(C) - log(current_popsize) << "\n";
		add += log(C) - log(current_popsize);
		toreturn += add;
		//cout << i << " "<< add << "\n";
		previous_coaltime = current_coaltime; //increment, save coalescent time
		i = i+1;
	}
	return toreturn;

}

double Tree::llk(int st_nlin, vector<double> coaltimes, vector<double>* bpoints, vector<double>* popsizes){

	// From Griffiths and Tavare (1994). Good summary in Heled and Drummond (2008)
	//
	// l(T| N(t) = \prod_{i=2}^n \frac{i choose 2}{N(t_i)} \exp {- \int_{t_{i+1}}^t_i \frac{i choose 2}{N(t)} dt}


	double toreturn = 0;
	int i = 0; //index coaltimes
	int j = 0; //index bpoints
	int k = 0; //index popsizes
	double previous_coaltime = 0;
	double previous_bpoint = 0;
	double previous_popsize = popsizes->at(0);
	while (i < coaltimes.size()){
		double current_coaltime = coaltimes.at(i);
		double current_bpoint = bpoints->at(j);
		double current_popsize = popsizes->at(k);
		int nlin = st_nlin-i+1;
		//cout << i << " "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
		double add = 0;
		double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		//cout << "C "<< C << "\n";
		while (current_bpoint < current_coaltime){



			//cout <<"here "<< i << " "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";

			//add the bit after the previous event
			if (previous_coaltime < current_bpoint && previous_coaltime > previous_bpoint)	{

				//cout << i << " here2 "<< current_coaltime << " "<< current_bpoint << " "<< previous_bpoint << "\n";
					//cout << i << " adding2 "<< - (current_bpoint - previous_coaltime) * C / current_popsize << "\n";
					//cout << previous_bpoint - previous_coaltime << " "<< C / previous_popsize << "\n";
					add += - (current_bpoint - previous_coaltime) * C / current_popsize;

			}
			else	{
				//cout << i << " here3 "<< current_bpoint << " "<< previous_bpoint << " "<< current_popsize << "\n";
				//cout << i  << " adding4 " << - (current_bpoint - previous_bpoint)* C / current_popsize <<"\n";
				add += - (current_bpoint - previous_bpoint)* C / current_popsize;
			}
			//increment
			j++;
			k++;
			previous_bpoint = current_bpoint;
			previous_popsize = current_popsize;
			current_bpoint = bpoints->at(j);
			current_popsize = popsizes->at(k);

		}
		if (previous_coaltime > previous_bpoint){
			//cout << "here4 "<< current_coaltime << " "<< previous_coaltime << " "<< C << " "<< current_popsize << " \n";
			//cout << i  << " adding5 " <<-(current_coaltime - previous_coaltime)*  C/ current_popsize << "\n";
			add += -(current_coaltime - previous_coaltime)*  C/ current_popsize;
		}
		else {
			//cout << i << " adding6 " << - (current_coaltime - previous_bpoint) * C /current_popsize << "\n";
			add += - (current_coaltime - previous_bpoint) * C /current_popsize;
		}
		//cout << i << " adding7 " << log(C) - log(current_popsize) << "\n";
		add += log(C) - log(current_popsize);
		toreturn += add;
		//cout << i << " "<< add << "\n";
		previous_coaltime = current_coaltime; //increment, save coalescent time
		i = i+1;
	}
	// now all modeled coalescents have happened
	// if there is still more than one lineage, need to add likelihood of *no* additional coalescences in the remaining time
	//
	if (st_nlin > coaltimes.size()+1){
		double add = 0;
		double lastcoal = coaltimes.at(coaltimes.size()-1);
		int nlin = st_nlin -coaltimes.size();
		double C = gsl_sf_fact(nlin) / (gsl_sf_fact(2)* gsl_sf_fact(nlin-2));
		for (int i = 0; i < bpoints->size(); i++){
			double dt;
			double N = popsizes->at(i);
			if (bpoints->at(i) < lastcoal) continue;
			if (i > 0 && bpoints->at(i) > lastcoal && bpoints->at(i-1) < lastcoal){
				dt = bpoints->at(i) - lastcoal;
			}
			else if (i > 0){
				dt = bpoints->at(i) - bpoints->at(i-1);
			}
			add +=  - dt * C/ N;
		}
		toreturn += add;
	}
	return toreturn;

}


void Tree::print(){
	cout << tree->get_newick_format() <<"\n";
}

void Tree::set_heights(){
	tree->set_heights();
}


void Tree::init_tip_labels(double h, gsl_rng * r){
	/*
	 * initialize labels at random with a given split time
	 *
	 */
	Graph::vertex_descriptor best;
	double besth;
	vector<Graph::vertex_descriptor> inorder = tree->get_inorder_traversal(npop);
	best = inorder[0];
	besth = 0;
	for (vector<Graph::vertex_descriptor>::iterator it = inorder.begin(); it != inorder.end(); it++){
		double testh = tree->g[*it].height;
		//if (tree->g[*it].index == 10) cout << "10 "<< tree->g[*it].height << "\n";
		//if (tree->g[*it].index == 0) cout << "0 "<< tree->g[*it].height << "\n";
		if (besth < h && testh > besth){
			besth = testh;
			best = *it;
		}
		else if (besth > h && testh < besth && testh > h){
			besth = testh;
			best = *it;
		}
	}
	if (besth < h){
		//cout << "here\n";
		//cout << tree->g[best].index << "\n";
		//cout << tree->g[best].is_root << "\n";
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5) {
			assign_labels(best, 1);
			while(!tree->g[best].is_root){

				best = tree->get_parent_node(best).first;
				assign_labels(best, 1);
			}
		}
		else {
			assign_labels(best, 2);
			while(!tree->g[best].is_root){

				best = tree->get_parent_node(best).first;
				assign_labels(best, 2);
			}
		}
	}
	else{
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(best);
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5){
			assign_labels(ch.first, 1);
			assign_labels(ch.second, 2);
		}
		else{
			assign_labels(ch.first, 2);
			assign_labels(ch.second, 1);

		}
		while (!tree->g[best].is_root){
			assign_labels(best, r, h);
			best = tree->get_parent_node(best).first;
		}
		//cout << tree->g[best].index << "\n";
		assign_labels(best, r, h);
	}
}

void Tree::assign_labels(Graph::vertex_descriptor v, int label){
	tree->g[v].popid = label;
	if (!tree->g[v].is_tip){
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(v);
		assign_labels(ch.first, label);
		assign_labels(ch.second, label);
	}
}


void Tree::assign_labels(Graph::vertex_descriptor v, gsl_rng* r, double h){
	if (tree->g[v].height > h ){
		tree->g[v].popid = 3;
		pair<Graph::vertex_descriptor, Graph::vertex_descriptor> ch = tree->get_child_nodes(v);
		//cout << tree->g[ch.first].index << " "<< tree->g[ch.first].popid << " "<< tree->g[ch.second].index << " "<< tree->g[ch.second].popid << "\n";
		if (tree->g[ch.first].popid < 1) {
			//cout << tree->g[v].index << " "<< tree->g[ch.first].index << "\n";
			assign_labels(ch.first, r, h);
		}
		if (tree->g[ch.second].popid < 1) {
			//cout << tree->g[v].index << " "<< tree->g[ch.first].index << "\n";
			assign_labels(ch.second, r, h);
		}
	}
	else{
		double ran = gsl_rng_uniform(r);
		if (ran < 0.5) assign_labels(v, 1);
		else assign_labels(v, 2);
	}
}

